//
// Created by zhou on 2021/10/26.
//
#include "yolov5/yolov5.h"
#include "arcface/arcface.h"
#include <ros/ros.h>
#include <image_transport/image_transport.h>
#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>
#include <drive/FaceId.h>
#include <vector>
#include "yolov5/dataset.h"
#include "data/results.h"
#include "arcface/draw_mask.h"
#include <drive/BoundingBox.h>
#include <drive/BoundingBoxes.h>


using namespace std;
using namespace std::chrono::_V2;
static unsigned int DoFaceIDTimes=0;
DriverIDResult CurDriverIDResult;
bool FaceCaptured=false;
drive::BoundingBox faceBox;
ros::Time  LatestResultStamp;
float tmpFaceFeature[ArcFace::FACE_FEATURE_DIMENSION]{0.0};
const float IOU_FACE_RECT_THRD=0.1;
ros::Time CurRecogStamp;
ros::Publisher StampInfoPub;
const unsigned int FACE_ID_TIME_PERIOD=1;
bool IsValidDriver=false;
static bool isCapFace=false;
string DriverName="";
void doFaceID(cv::Mat &frameDriver, ros::Time stampDriver,cv::Rect faceRect)
{
    // �±�����Խ�磬���뱣��һ�£���Ȼ���쳣����
    if (faceRect.x > 0 && faceRect.y > 0 && faceRect.width > 0 && faceRect.height > 0 &&
        (faceRect.x + faceRect.width) < frameDriver.cols &&
        (faceRect.y + faceRect.height) < frameDriver.rows)
    {
        cv::Mat faceMat = frameDriver(faceRect);
        // cout << "����������ȷλ�ã���ʼ������ʻԱ����..." << endl;
        CurRecogStamp = stampDriver;
        int faceId=0;
        ArcFace::DetectFaceID(faceMat, faceId, tmpFaceFeature);
        CurDriverIDResult.AddOneRecogResult( faceId, tmpFaceFeature);
        cout << "��������-----------------------------------------ID: " << faceId << endl << endl;
    }
}
void doDriverID(cv::Mat &frameDriver, ros::Time &stampDriver,cv::Rect faceRect)
{
    cv::Rect maskFaceRect = ArcFace::GetFacePos(frameDriver.cols, frameDriver.rows);
    cout << "faceRect: " << faceRect.x << ", " << faceRect.y << ", " << faceRect.width << ", " << faceRect.height << endl;
    cout << "maskFaceRect: " << maskFaceRect.x << ", " << maskFaceRect.y << ", " << maskFaceRect.width << ", " << maskFaceRect.height << endl;
    cout << "IOU: " << ResultTool::IOU(faceRect, maskFaceRect) << endl;

    if (ResultTool::IOU(faceRect, maskFaceRect) >= IOU_FACE_RECT_THRD)
    {
        isCapFace = true;
        doFaceID(frameDriver, stampDriver, faceRect);
        CurDriverIDResult.GetIDReuslt();
        // ��������ʶ��������Σ�������ȷ���ˣ�Ȼ������������ҵ����
        if(CurDriverIDResult.ID>0&&CurDriverIDResult.ID<2000)
        {
            IsValidDriver=true;
            cout << "the driver -----------------------------------------ID: " <<CurDriverIDResult.ID << endl ;
            DriverName=CurDriverIDResult.QueryDriverName();
            
            cout<<"the driver name is:"<<DriverName<<endl;
        }
        else {
            // IsValidDriver=false;
            cout << "the driver-----------------------------------------ID: " <<CurDriverIDResult.ID << endl ;
            // for(auto iter=CurDriverIDResult.DriverIds.begin();iter!=CurDriverIDResult.DriverIds.end();iter++){
            //     cout<<*iter<<",";
            // }
        }
        DoFaceIDTimes++;
        if (DoFaceIDTimes > DriverIDResult::FACE_RECOG_TIMES)
        {
            CurDriverIDResult.GetIDReuslt();
            DoFaceIDTimes = 0;
        }
    }
    else
    {
        isCapFace = false;
        return;
    }
}
void MsgCallback(const drive::BoundingBoxes::ConstPtr& msg){
    if(msg->FaceCaptured){
        FaceCaptured=true;
        LatestResultStamp=msg->LatestResultStamp;
        for(auto iter=msg->boundingboxes.begin();iter!=msg->boundingboxes.end();iter++){
            if(iter->Class=="FACE"){
                faceBox.Class=iter->Class;
                faceBox.Rect_x=iter->Rect_x;
                faceBox.Rect_y=iter->Rect_y;
                faceBox.Rect_w=iter->Rect_w;
                faceBox.Rect_h=iter->Rect_h;
            }
        }
    }
    else{
        FaceCaptured=false;
    }
}
void imageCallback(const sensor_msgs::ImageConstPtr& msg){
    try {
        ros::Time nowStamp = ros::Time::now();//����ʱ�����һ��ʱ����ǵ�ǰ��ʱ�䣬��һ��ʱ����ǵ�ǰ������ͼƬ������ʱ���
        ros::Time stamp = msg->header.stamp;
        cout << "the duration is  " << (int) ((nowStamp - stamp).toSec() * 1000) << " ms" << endl;

        cv::Mat frame = cv_bridge::toCvShare(msg, "bgr8")->image;
        drive::FaceId face_msg;//����һ��msg����ý������Է�����ȥ
        face_msg.CurFaceRecogStamp = CurRecogStamp;
        face_msg.IsValidDriver = IsValidDriver;
        face_msg.DriverName=DriverName;
        StampInfoPub.publish(face_msg);
        if(IsValidDriver==false)
        {
            cout<<"there isn't valid driver!"<<endl;
            // if((int) (CurRecogStamp - stamp).toSec() > FACE_ID_TIME_PERIOD){
                if(FaceCaptured){
                    cout<<"facecaptured!"<<endl;
                    // float bbox[4];
                    // bbox[0]=faceBox.Rect_x;
                    // bbox[1]=faceBox.Rect_y;
                    // bbox[2]=faceBox.Rect_w;
                    // bbox[3]=faceBox.Rect_h;
                    // cout<<"bbox="<<bbox[0]<<","<<bbox[1]<<endl;
                    cv::Rect faceRect;
                    faceRect.x=faceBox.Rect_x;
                    faceRect.y=faceBox.Rect_y;
                    faceRect.width=faceBox.Rect_w;
                    faceRect.height=faceBox.Rect_h;
                     if (faceRect.x > 0 && faceRect.y > 0 && faceRect.width > 0 && faceRect.height > 0 &&
                    (faceRect.x + faceRect.width) < frame.cols &&
                    (faceRect.y + faceRect.height) < frame.rows) {
                    //cv::Mat frameface = frame(faceRect);
                    cout<<"do driver id!"<<endl;
                    doDriverID(frame,stamp,faceRect);
                }
                // }
            }
        }
        
        // if ((int) (CurRecogStamp - stamp).toSec() > FACE_ID_TIME_PERIOD) {
        //     face_msg.CurFaceRecogStamp = CurRecogStamp;
        //     face_msg.IsValidDriver = IsValidDriver;
        //     StampInfoPub.publish(face_msg);


        // }
        // vector <Yolo::Detection> result = YoloV5::AnalyzeOneShot(frame);
        // for (auto iter = result.begin(); iter != result.end(); iter++) {
        //     if (iter->class_id == ClassID::FACE) {
        //         cv::Rect faceRect = YoloV5::get_rect(frame, iter->bbox);
        //         if (faceRect.x > 0 && faceRect.y > 0 && faceRect.width > 0 && faceRect.height > 0 &&
        //             (faceRect.x + faceRect.width) < frame.cols &&
        //             (faceRect.y + faceRect.height) < frame.rows) {
        //             //cv::Mat frameface = frame(faceRect);
        //             doDriverID(frame,stamp,faceRect);
        //         }
        //     }
        // }
    }
    catch (cv_bridge::Exception& e)
    {
        ROS_ERROR("Could not convert from '%s' to 'bgr8'.", msg->encoding.c_str());
    }
}
int main(int argc, char **argv)
{
    ros::init(argc, argv, "image_faceid_node");
    ros::NodeHandle node_faceid;

    //��ֹ����������ģ���������ȳ�ʼ��
    cudaSetDevice(DEVICE);
    ArcFace::InitArcFaceEngine();
    cout<<"ArcFace引擎序列化完成ing"<<endl;

    ArcFace::ReadFaceDataToGPU();
    cout<<"loading  face data successfully!"<<endl;
    // YoloV5::InitYoloV5Engine();
    // cout << "faceid YoloV5 �������л����" << endl;
    CurRecogStamp = ros::Time::now();
    StampInfoPub = node_faceid.advertise<drive::FaceId>("/camera_csi0/cur_faceid_result", 1);

    image_transport::ImageTransport it(node_faceid);
    image_transport::Subscriber sub = it.subscribe("/camera_csi0/frames", 1, imageCallback);
    ros::Subscriber person_info_sub = node_faceid.subscribe("/camera_csi0/cur_result",1,MsgCallback);
    ros::spin();

    // YoloV5::ReleaseYoloV5Engine();
    ArcFace::ReleasePFLDEngine();
}
